import warnings

warnings.filterwarnings('ignore')
from ultralytics import YOLO

data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\DDD_YOLO.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_721.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_403525.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_unresize.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\TIB-NET.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_B0.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_B10.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_B1.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_B1_5000.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_B3.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_B3_5000.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_B7.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_B7_5000.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_B9.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_B9_5000.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_B5_5000.yaml'
# data_cfg = r'E:\pycharmProjs\yolo-ladcs\dataset\SDB9K_B5_5000_2.yaml'

model_scale = 'n'

if __name__ == '__main__':
    # model = YOLO('ultralytics/cfg/models/v8/yolov8n.yaml')
    # model = YOLO('ultralytics/cfg/models/v8/yolov8n-ladc.yaml')
    # model = YOLO('ultralytics/cfg/models/v8/yolov8n-ladcs.yaml')
    model = YOLO(f'ultralytics/cfg/models/v8/yolov8{model_scale}.yaml')
    model.load(f'yolov8{model_scale}.pt')  # loading pretrain weights
    # model.load('yolov8s.pt')  # loading pretrain weights
    model.train(data=data_cfg,
                cache=False,
                # imgsz=640,
                imgsz=640,
                epochs=150,
                batch=64,
                close_mosaic=10,
                workers=4,
                device='0',
                optimizer='SGD',  # using SGD
                # optimizer='AdamW', # using AdamW
                # resume=r'E:\pycharmProjs\yolo-ladcs\runs\train\exp17\weights\last.pt', # last.pt path
                # amp=False, # close amp
                # fraction=0.2,
                project='runs/train',
                name='exp',
                )
